from fmpy import dump, simulate_fmu
import numpy as np

from fmpy import dump, simulate_fmu
import numpy as np
import os, time, warnings

if os.path.exists(r"./DishWasherROM.fmu"):
    FMU = r"./DishWasherROM.fmu"
else:
    FMU = r"./lib/DishWasherROM.fmu"
    
## IMPORT THE FOLLOWING JUST FOR DATA GENERATION
from multiprocessing import Process, Queue

dump(FMU)
TARGET_ANGLE = np.pi/2
SAMPLE_DT = 0.05; TS = 10
TIMEREF = np.round(np.arange(round(TS/SAMPLE_DT)+1)*SAMPLE_DT, int(-np.log(SAMPLE_DT))+1)
DOOR_ANGLE = TIMEREF * TARGET_ANGLE / TS
DTYPE = [('time', "float32"), ('Door_Ang', "float32"), ('JL_Damage', "float32"), ('LS_DAMAGE', "float32"),
         ('MFT_DAMAGE', "float32"), ('SK_Damage', "float32"), ('TH_Damage', "float32")]

def simulator(theta):
    #t = time()
    X = np.array([(time_r, door_ang, *theta)
                  for time_r, door_ang in zip(TIMEREF, DOOR_ANGLE)], dtype=DTYPE)
    res = simulate_fmu(FMU, record_events=False, validate=False,
                       input=X, stop_time=TS, step_size=SAMPLE_DT).tolist()
    #print(time()-t)
    return np.concatenate([np.array(X.tolist(), dtype="float32")[:, 1:], np.array(res, dtype="float32")[:, 1:]], axis=1)

import matplotlib.pyplot as plt

def _get_params(theta=[0,0,0,0,0], plot_res=False, r_p=2e-5):
    x = simulator(theta)
    reg_params = {"ruls": [float(r_theta) for r_theta in theta]}
    inds = [6,7,8,11]
    inds_pre = [6,7,8,9,10,11]
    name = ["Door_Force", "JL_Fz", "JL_Fy", "MFT_Force", "SK_Force", "TH_Force"]
    
    if plot_res:
        fig = plt.figure()
        axs = [fig.add_subplot(3,2,_id+1) for _id in range(6)]
        for i in range(6):
            axs[i].plot(x[:,0], x[:,6+i])

    x_ext_all = np.ones((len(x), 7))
    x_ext_all[:,0]=x[:,0]
    x_ext_all[:,1]=np.sin(x[:,0])
    x_ext_all[:,2]=np.cos(x[:,0])
    x_ext_all[:,3]=x[:,0]*np.sin(x[:,0])
    x_ext_all[:,4]=x[:,0]*np.cos(x[:,0])
    x_ext_all[:,5]=x[:,0]**2

    x_= x[x[:,0]<0.05]
    x_ext = np.ones((len(x_), 7))
    x_ext[:,0]=x_[:,0]
    x_ext[:,1]=np.sin(x_[:,0])
    x_ext[:,2]=np.cos(x_[:,0])
    x_ext[:,3]=x_[:,0]*np.sin(x_[:,0])
    x_ext[:,4]=x_[:,0]*np.cos(x_[:,0])
    x_ext[:,5]=x_[:,0]**2

    thres_ind = np.where(x[:,9] > 0.02)[0][-1]
    print("SPLIT IND:", thres_ind)
    reg_params["split_ind"] = int(thres_ind)

    print()
    
    for i in inds_pre:
        y_ = x_ [:,i,None]
        params = np.linalg.pinv(x_ext) @ y_
        params[np.abs(params)*np.abs(x_ext).max(axis=0)[:,None]<np.abs(params*np.abs(x_ext).max(axis=0)[:,None]).max()*r_p] = 0
        reg_params[f"{name[i-6]}_pre"] = params.flatten().astype(float).tolist()
        err = np.mean((y_ - x_ext @ params)**2)**0.5
        print(f"{name[i-6]} = {params[0,0]}*alpha + {params[1,0]}*sin(alpha) + {params[2,0]}*cos(alpha) + {params[3,0]}*alpha*sin(alpha) + {params[4,0]}*alpha*cos(alpha) + {params[5,0]}*alpha**2 +{params[-1,0]}")
        print(f"Absolute rr = {err}\t Relative Err = {err/np.std(y_)}")
        if plot_res:
            axs[i-6].plot(x[:thres_ind, 0], (x_ext_all[:thres_ind] @ params).flatten(), "y--")        

    print()

    x_= x[thres_ind:]
    x_ext = np.ones((len(x_), 7))
    x_ext[:,0]=x_[:,0]
    x_ext[:,1]=np.sin(x_[:,0])
    x_ext[:,2]=np.cos(x_[:,0])
    x_ext[:,3]=x_[:,0]*np.sin(x_[:,0])
    x_ext[:,4]=x_[:,0]*np.cos(x_[:,0])
    x_ext[:,5]=x_[:,0]**2
    
    for i in inds_pre:
        if i in inds:
            y_ = x_ [:,i,None]
            params = np.linalg.pinv(x_ext) @ y_
            params[np.abs(params)*np.abs(x_ext).max(axis=0)[:,None]<np.abs(params*np.abs(x_ext).max(axis=0)[:,None]).max()*r_p] = 0
            reg_params[f"{name[i-6]}_post"] = params.flatten().astype(float).tolist()
            err = np.mean((y_ - x_ext @ params)**2)**0.5
            print(f"{name[i-6]} = {params[0,0]}*alpha + {params[1,0]}*sin(alpha) + {params[2,0]}*cos(alpha)  + {params[3,0]}*alpha*sin(alpha) + {params[4,0]}*alpha*cos(alpha) + {params[5,0]}*alpha**2 +{params[-1,0]}")
            print(f"Absolute rr = {err}\t Relative Err = {err/np.std(y_)}")
            if plot_res:
                axs[i-6].plot(x[thres_ind:, 0], (x_ext_all[thres_ind:] @ params).flatten(), "g--")
        else:
            y_ = x_ [:,i,None]
            err = np.mean(y_**2)**0.5
            print(f"{name[i-6]} = 0")
            print(f"Absolute rr = {err}\t Relative Err = {err/np.std(y_) if np.std(y_) < 1e-10 else err}")
            if plot_res:
                axs[i-6].plot(x[thres_ind:, 0], [0 for _ in x[thres_ind:, 0]], "y--")

    if plot_res:
        plt.show()
    
    return reg_params

def _random_gene(_theta, n_in, q,):
    _params = _get_params(theta=_theta, plot_res=False)
    q.put([_params, n_in])

def batch_generate(x=None, worker_num=5, verbose=False, filter_na=True):
    n = len(x)
    ys = [None for _ in range(n)]
    
    q = Queue(); n_in = 0
    workers_on = 0
    Workers = [None for _ in range(worker_num)]
    
    try:
        while n_in < n or workers_on > 0:
            for w in range(worker_num):
                if Workers[w] is None:
                    if n_in < n:
                        Workers[w] = Process(target=_random_gene, args=(None if x is None or n_in >= len(x) else x[n_in], n_in, q,))
                        n_in += 1
                        if verbose and (n_in + 1) % 10 == 0:
                            print(f"Samples {n_in+1}/{n}")
                        if Workers[w]:
                            Workers[w].start()
                            workers_on += 1

                elif not Workers[w].is_alive():
                    if n_in < n:
                        Workers[w].join()
                        workers_on -= 1
                        Workers[w] = Process(target=_random_gene, args=(None if x is None or n_in >= len(x) else x[n_in], n_in, q,))
                        n_in += 1
                        if verbose and (n_in + 1) % 10 == 0:
                            print(f"Samples {n_in+1}/{n}")
                        if Workers[w]:
                            Workers[w].start()
                            workers_on += 1
                    else:
                        Workers[w].join()
                        workers_on -= 1
            while not q.empty():
                y_, n_in_ = q.get_nowait()
                if verbose:
                    print(f"\t Sample {n_in_+1} is registered")
                ys[n_in_] = y_
            time.sleep(0.1)
        while not q.empty():
            y_, n_in_ = q.get_nowait()
            ys[n_in_] = y_
        if verbose:
            print("[Success] Data is already Generated")
    except Exception as e:
        raise Exception("[Error]",e)
    if filter_na:
        return [itm for itm in ys if not itm is None][:n]
    else:
        return ys

def _convert_param(param):
    flat_param = [param["split_ind"]]
    name = ["Door_Force", "JL_Fz", "JL_Fy", "MFT_Force", "SK_Force", "TH_Force"]
    pname = [f"{n}_{n_pre}" for n in name for n_pre in ["pre", "post"] if f"{n}_{n_pre}" in param.keys()]
    for pn in pname:
        flat_param.extend(param[pn])
    return flat_param

def _pinv(x):
    d,v = np.linalg.eig(x.T @ x)
    xinv = (v @ np.diag([0 if abs(d_) <= 1e-5 else 1/d_ for d_ in d]) @ np.linalg.pinv(v)) @ x.T
    return xinv
    
def _deep_params(params, r_p=1e-10, n_eps=10):
    pLen = len(params[0]["MFT_Force_pre"])
    xs, ys = zip(*[[param["ruls"]+[1], _convert_param(param)] for param in params])
    xs = np.array(xs, dtype="float32")
    xs_ext = np.concatenate([(x_a*x_b)[:,None]#*x_c*x_d*x_e*x_f*x_g)[:,None]
                             for x_i, x_a in enumerate(xs.T[:])
                             for x_j, x_b in enumerate(xs.T[x_i:])
                             #for x_k, x_c in enumerate(xs.T[x_i+x_j:])
                             #for x_l, x_d in enumerate(xs.T[x_i+x_j+x_k:])
                             #for x_m, x_e in enumerate(xs.T[x_i+x_j+x_k+x_l:])
                             #for x_n, x_f in enumerate(xs.T[x_i+x_j+x_k+x_l+x_m:])
                             #for x_o, x_g in enumerate(xs.T[x_i+x_j+x_k+x_l+x_m+x_n:])
                             ],axis=1)
    ys = np.array(ys, dtype="float32")
    param = np.round((_pinv(xs_ext) @ ys), n_eps)
    print(ys.shape, param.shape)
    param[np.abs(param)*np.abs(xs_ext).max(axis=0)[:,None]<np.abs(param*np.abs(xs_ext).max(axis=0)[:,None]).max()*r_p] = 0
    print("Relative Err:", np.mean((xs_ext @ param - ys)**2/np.std(ys, axis=0)**2)**0.5)
    param = param.T
    param_ext = param[1:].reshape((-1, pLen, param.shape[-1]))
    return [param[0].astype(float).tolist()] + param_ext.astype(float).tolist()

def _inner_formula(param):
    pns_ = ["JL_Damage", "LS_DAMAGE", "MFT_DAMAGE", "SK_Damage", "TH_Damage", "1"]
    pns = [f"{px_a}*{px_b}"#*{px_c}*{px_d}*{px_e}*{px_f}*{px_g}"
           for px_i, px_a in enumerate(pns_[:])
           for px_j, px_b in enumerate(pns_[px_i:])
           #for px_k, px_c in enumerate(pns_[px_i+px_j:])
           #for px_l, px_d in enumerate(pns_[px_i+px_j+px_k:])
           #for px_m, px_e in enumerate(pns_[px_i+px_j+px_k+px_l:])
           #for px_n, px_f in enumerate(pns_[px_i+px_j+px_k+px_l+px_m:])
           #for px_o, px_g in enumerate(pns_[px_i+px_j+px_k+px_l+px_m+px_n:])
           ]
    inner_code = "+".join([f"{p}*{pn}" for p, pn in zip(param, pns)])
    return inner_code.replace("+-", "-")    

def _formula(params, pnames):
    code = "import numpy as np\nfrom numpy import sin,cos\ndef _fmu_approx(alpha, theta):\n    "
    split_ind = "0.007854*(" + _inner_formula(params[0]) + ")"
    pns_ = ["JL_Damage", "LS_DAMAGE", "MFT_DAMAGE", "SK_Damage", "TH_Damage"]
    code += f"{','.join(pns_)}=theta\n    alpha_0 = {split_ind}\n    "
    for pms, pns in zip(params[1:], pnames):
        pn_ = pns[::-1].split("_", 1)[1][::-1]
        pm_code = [_inner_formula(pm) for pm in pms]
        inner_pm_c = "(" + "+".join([f"({pm_c})*{pn_c}" for pm_c, pn_c in zip(pm_code, ["alpha", "sin(alpha)", "cos(alpha)", "alpha*sin(alpha)", "alpha*cos(alpha)", "alpha**2"])]) + f"+{pm_code[-1]})"
        inner_pm_c = inner_pm_c.replace("+-", "-")
        if "_pre" in pns:
            code += f"{pn_}=np.zeros(len(alpha))\n    {pn_}_ = {inner_pm_c}\n    {pn_}[alpha<alpha_0]={pn_}_[alpha<alpha_0]\n    "
        else:
            code += f"{pn_}_ = {inner_pm_c}\n    {pn_}[alpha>=alpha_0]={pn_}_[alpha>=alpha_0]\n    "
    return code +"""return Door_Force, JL_Fz, JL_Fy, MFT_Force, SK_Force, TH_Force

TARGET_ANGLE = np.pi/2
SAMPLE_DT = 0.05; TS = 10
TIMEREF = np.round(np.arange(round(TS/SAMPLE_DT)+1)*SAMPLE_DT, int(-np.log(SAMPLE_DT))+1)
DOOR_ANGLE = TIMEREF * TARGET_ANGLE / TS

def simulator_approx(theta):
    return np.concatenate([itm[:,None] for itm in _fmu_approx(DOOR_ANGLE, theta)], axis=1)

if __name__ == "__main__":
    theta = [0,0,0,0,0]
    res = simulator_approx(theta)

    from fmpy import dump, simulate_fmu
    FMU = r"./DishWasherROM.fmu"
    dump(FMU)
    DTYPE = [('time', "float32"), ('Door_Ang', "float32"), ('JL_Damage', "float32"), ('LS_DAMAGE', "float32"),
         ('MFT_DAMAGE', "float32"), ('SK_Damage', "float32"), ('TH_Damage', "float32")]
    
    res_gnd = simulate_fmu(FMU, record_events=False, validate=False,
                       input=np.array([(time_r, door_ang, *theta)
                                       for time_r, door_ang in zip(TIMEREF, DOOR_ANGLE)], dtype=DTYPE),
                       stop_time=TS, step_size=SAMPLE_DT).tolist()
    res_gnd = np.array(res_gnd, dtype="float32")[:,1:]
    
    import matplotlib.pyplot as plt
    fig = plt.figure()
    for res_ind, res_itm in enumerate(res.T):
        ax = fig.add_subplot(3,2, res_ind+1)
        ax.plot(res_itm)
        ax.plot(res_gnd[:,res_ind], "--")
    plt.show()"""

if __name__ == "__main__":
    _get_params(theta=[0, 0, 0, 0, 0], plot_res=True)

    import json
    regene = False
    
    if regene:
        N = 8
        xs = np.array([[i1/N, i2/N, i3/N, i4/N, i5/N] for i1 in range(N+1) for i2 in range(N+1)
                       for i3 in range(N+1) for i4 in range(N+1) for i5 in range(N+1)], dtype="float32")
        res_y = batch_generate(xs, verbose=True)
        with open(f"data_punch.json", "w+", encoding="gbk") as f:
            json.dump(res_y, f)
    else:
        with open(f"data_punch.json", "r+", encoding="gbk") as f:
            res_y = json.load(f)

    name = ["Door_Force", "JL_Fz", "JL_Fy", "MFT_Force", "SK_Force", "TH_Force"]
    pname = [f"{n}_{n_pre}" for n in name for n_pre in ["pre", "post"] if f"{n}_{n_pre}" in res_y[0].keys()]

    param = _deep_params(res_y)
    code = _formula(param, pname)

    with open("simu_approx.py", "w+") as f:
        f.write(code)
        
    """
    xs, ys = zip(*[[param["ruls"]+[1], _convert_param(param)] for param in res_y])
    xs = np.array(xs, dtype="float32")
    ys = np.array(ys, dtype="float32")#[ind]
    
    def _plot(k):
        fig = plt.figure()
        for i in range(5):
            ax = fig.add_subplot(3,2,i+1)
            _ind = [l for l in range(5) if l!=i]
            _x = (xs[:,_ind]*(11**np.arange(len(_ind)))).sum(axis=1)
            for xi in np.sort(np.unique(_x))[::100]:
                ax.plot(xs[:,i][_x==xi], ys[:,k][_x==xi], linewidth=0.5 if xi else 3)
            ax.set_ylabel(["JL_Damage", "LS_DAMAGE", "MFT_DAMAGE", "SK_Damage", "TH_Damage"][i])
        ax = fig.add_subplot(3,2,6)
        _ind = [l for l in range(5) if l not in [2,3]]
        _x = (xs[:,_ind]*(11**np.arange(len(_ind)))).sum(axis=1)
        for xi in np.sort(np.unique(_x))[::10]:
            xs_ = (np.minimum(0.8, xs[:,2])*np.minimum(0.8, xs[:,3]))[_x==xi]
            ys_ = ys[:,k][_x==xi]
            ind = np.argsort(xs_)
            ax.plot(xs_[ind],ys_[ind], linewidth=0.5 if xi else 3)
            
        fig.tight_layout()
		
    for il in range(5):
        _plot(il)
    """
    
    plt.show()
